import cv2
import numpy as np
from typing import List, Tuple

class MapGenerator:
    def __init__(self, resolution: float = 0.05, initial_buffer: float = 10.0):
        """
        初始化地图生成器
        resolution: 地图分辨率 (米/像素)
        initial_buffer: 初始地图边缘预留空间 (米)
        """
        self.resolution = resolution
        self.initial_buffer = initial_buffer
        
        # 地图数据
        self.map_image = None  # 地图图像
        self.origin = (0.0, 0.0)  # 地图原点 (x, y)
        self.width = 0  # 地图宽度 (像素)
        self.height = 0  # 地图高度 (像素)
        
        # 记录的点
        self.points = []  # 存储为像素坐标
        self.world_points = []  # 存储为世界坐标
        
        # 机器人状态标志
        self.robot_flag = 0
        
        # 自适应原点相关
        self.min_x = float('inf')
        self.min_y = float('inf')
        self.max_x = -float('inf')
        self.max_y = -float('inf')
        
    def update_bounds(self, x: float, y: float):
        """更新坐标边界"""
        self.min_x = min(self.min_x, x)
        self.min_y = min(self.min_y, y)
        self.max_x = max(self.max_x, x)
        self.max_y = max(self.max_y, y)
        
    def initialize_map(self):
        """根据当前边界初始化地图"""
        # 计算地图尺寸
        width = int((self.max_x - self.min_x) / self.resolution) + 1
        height = int((self.max_y - self.min_y) / self.resolution) + 1
        
        # 添加初始缓冲区
        width += int(self.initial_buffer / self.resolution) * 2
        height += int(self.initial_buffer / self.resolution) * 2
        
        # 计算新原点
        origin_x = self.min_x - self.initial_buffer
        origin_y = self.min_y - self.initial_buffer
        
        # 创建新地图
        self.map_image = np.ones((height, width), dtype=np.int8) * 255
        self.origin = (origin_x, origin_y)
        self.width = width
        self.height = height
        
        # 将已记录的点转换到新地图坐标系
        self.points = []
        for world_point in self.world_points:
            pixel_x, pixel_y = self.world_to_pixel(world_point[0], world_point[1])
            self.points.append((pixel_x, pixel_y))
            
        # 如果有点，更新地图
        if len(self.points) >= 3:
            self.update_map()
            
        return True
        
    def world_to_pixel(self, x: float, y: float) -> Tuple[int, int]:
        """将世界坐标转换为像素坐标"""
        pixel_x = int((x - self.origin[0]) / self.resolution)
        pixel_y = int((y - self.origin[1]) / self.resolution)
        return pixel_x, pixel_y
        
    def add_point(self, x: float, y: float):
        """添加一个点到地图"""
        # 更新边界
        self.update_bounds(x, y)
        
        # 存储世界坐标
        self.world_points.append((x, y))
        
        # 如果地图未初始化，检查是否可以初始化
        if self.map_image is None:
            # 检查是否有足够的点来确定边界
            if len(self.world_points) >= 2:
                return self.initialize_map()
            return False
            
        # 转换为像素坐标
        pixel_x, pixel_y = self.world_to_pixel(x, y)
        
        # 存储像素坐标
        self.points.append((pixel_x, pixel_y))
        
        # 如果有足够的点形成闭合区域，更新地图
        if len(self.points) >= 3:
            self.update_map()
            
        return True
        
    def set_robot_flag(self, flag: int):
        """设置机器人状态标志"""
        self.robot_flag = flag
        
    def update_map(self):
        """根据记录的点更新地图"""
        if len(self.points) < 3:
            return
            
        # 创建一个临时图像用于绘制多边形
        temp_image = np.zeros((self.height, self.width), dtype=np.uint8)
        
        # 将点列表转换为NumPy数组
        pts = np.array(self.points, np.int32)
        pts = pts.reshape((-1, 1, 2))
        
        # 在临时图像上绘制填充多边形
        cv2.fillPoly(temp_image, [pts], 1)
        
        # 将多边形区域内的栅格值设为机器人状态标志
        mask = temp_image == 1
        self.map_image[mask] = self.robot_flag
        
    def get_map_data(self) -> Tuple[np.ndarray, Tuple[float, float], int, int, float]:
        """获取地图数据和相关参数"""
        return self.map_image, self.origin, self.width, self.height, self.resolution